"""
pinecone_train.py - 使用 Pinecone 创建索引并上传 MNIST 数据

要求：
1. 参考使用Pinecone进行数字判断的实例 借鉴代码实现
2. 用80%的mnist数据创建 Pinecone 的索引
3. 用20%的数据测试当 k=11 时准确率
4. 最终用 logging 打印：
   - 成功创建索引，并上传了 1437 条数据
   - 当 k=11 是，使用 Pinecone 的准确率
   - 上传数据和测试 k 的准确率的时候都要有进度条
   - 用 logging 打印的信息需要有日期

注意：要使用 py 文件，而不是 ipynb 文件
"""

import os
import time
import logging
from datetime import datetime
from typing import Tuple
import numpy as np
from sklearn import datasets
from sklearn.model_selection import train_test_split
from tqdm import tqdm

try:
    from pinecone import Pinecone, ServerlessSpec
    PINECONE_AVAILABLE = True
except ImportError:
    PINECONE_AVAILABLE = False
    print("错误: pinecone 未安装。请运行: pip install pinecone")
    exit(1)

# 配置 logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# Pinecone API Key 配置
PINECONE_API_KEY = "pcsk_6xFLgp_DBvvvh8Yzqz3qD3rH4KykRKBakV3Br9BtuQjUnccSNNuJ59BoaGzvLGtVRnocd9"
INDEX_NAME = "mnist-index"


def load_mnist_data() -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    加载 MNIST 数据集（使用 sklearn digits 作为替代）
    
    Returns:
        X_train, X_test, y_train, y_test
    """
    logger.info("正在加载 MNIST 数据集...")
    
    # 使用 sklearn digits 数据集（8x8 图像）
    digits = datasets.load_digits()
    X, y = digits.data, digits.target
    
    logger.info(f"数据集加载完成: {X.shape[0]} 个样本, 每个样本 {X.shape[1]} 维")
    
    # 划分训练集和测试集 (80% / 20%)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )
    
    logger.info(f"训练集: {X_train.shape[0]} 个样本 (80%)")
    logger.info(f"测试集: {X_test.shape[0]} 个样本 (20%)")
    
    return X_train, X_test, y_train, y_test


def create_pinecone_index(pc: Pinecone, index_name: str, dimension: int = 64):
    """
    创建 Pinecone 索引
    
    Args:
        pc: Pinecone 客户端
        index_name: 索引名称
        dimension: 向量维度
    """
    logger.info(f"正在检查索引 '{index_name}' 是否存在...")
    
    # 检查索引是否已存在
    existing_indexes = [idx.name for idx in pc.list_indexes()]
    
    if index_name in existing_indexes:
        logger.info(f"索引 '{index_name}' 已存在，正在删除...")
        pc.delete_index(index_name)
        time.sleep(2)
        logger.info(f"索引 '{index_name}' 已删除")
    
    # 创建新索引
    logger.info(f"正在创建索引 '{index_name}'...")
    pc.create_index(
        name=index_name,
        dimension=dimension,
        metric="euclidean",
        spec=ServerlessSpec(
            cloud='aws',
            region='us-east-1'
        )
    )
    
    # 等待索引准备就绪
    logger.info("等待索引准备就绪...")
    while not pc.describe_index(index_name).status['ready']:
        time.sleep(1)
    
    logger.info(f"索引 '{index_name}' 创建成功！")


def upload_data_to_pinecone(index, X_train: np.ndarray, y_train: np.ndarray, batch_size: int = 100):
    """
    上传训练数据到 Pinecone（带进度条）
    
    Args:
        index: Pinecone 索引对象
        X_train: 训练特征
        y_train: 训练标签
        batch_size: 批量大小
    """
    n_samples = X_train.shape[0]
    logger.info(f"开始上传 {n_samples} 条数据到 Pinecone...")
    
    # 使用 tqdm 显示进度条
    with tqdm(total=n_samples, desc="上传数据到Pinecone", unit="条") as pbar:
        for i in range(0, n_samples, batch_size):
            batch_end = min(i + batch_size, n_samples)
            batch_vectors = []
            
            for j in range(i, batch_end):
                vector_id = str(j)
                vector = X_train[j].tolist()
                label = int(y_train[j])
                
                batch_vectors.append({
                    "id": vector_id,
                    "values": vector,
                    "metadata": {"label": label}
                })
            
            # 上传批次
            index.upsert(vectors=batch_vectors)
            pbar.update(batch_end - i)
    
    logger.info(f"成功创建索引，并上传了 {n_samples} 条数据")
    
    # 等待索引更新
    time.sleep(2)


def test_accuracy_with_k(index, X_test: np.ndarray, y_test: np.ndarray, k: int = 11) -> float:
    """
    测试指定 k 值的准确率（带进度条）
    
    Args:
        index: Pinecone 索引对象
        X_test: 测试特征
        y_test: 测试标签
        k: KNN 的 k 值
        
    Returns:
        准确率
    """
    n_test = len(X_test)
    correct = 0
    
    logger.info(f"开始测试 k={k} 时的准确率...")
    
    # 使用 tqdm 显示进度条
    with tqdm(total=n_test, desc=f"测试 k={k} 的准确率", unit="个") as pbar:
        for i in range(n_test):
            # 查询 Pinecone
            query_vector = X_test[i].tolist()
            results = index.query(
                vector=query_vector,
                top_k=k,
                include_metadata=True
            )
            
            # 获取最近邻的标签
            labels = [match['metadata']['label'] for match in results['matches']]
            
            # 投票决定预测标签
            predicted_label = int(np.bincount(labels).argmax())
            
            # 检查是否正确
            if predicted_label == y_test[i]:
                correct += 1
            
            pbar.update(1)
    
    accuracy = correct / n_test
    logger.info(f"当 k={k} 时，使用 Pinecone 的准确率: {accuracy:.4f} ({correct}/{n_test})")
    
    return accuracy


def main():
    """主函数"""
    logger.info("="*60)
    logger.info("Pinecone MNIST 训练程序")
    logger.info("="*60)
    
    # 1. 加载数据
    X_train, X_test, y_train, y_test = load_mnist_data()
    
    # 2. 初始化 Pinecone
    logger.info("正在初始化 Pinecone...")
    try:
        pc = Pinecone(api_key=PINECONE_API_KEY)
        logger.info("Pinecone 初始化成功")
    except Exception as e:
        logger.error(f"Pinecone 初始化失败: {e}")
        return
    
    # 3. 创建索引
    create_pinecone_index(pc, INDEX_NAME, dimension=64)
    
    # 4. 连接到索引
    logger.info(f"正在连接到索引 '{INDEX_NAME}'...")
    index = pc.Index(INDEX_NAME)
    logger.info("索引连接成功")
    
    # 5. 上传训练数据
    upload_data_to_pinecone(index, X_train, y_train)
    
    # 6. 测试准确率 (k=11)
    accuracy = test_accuracy_with_k(index, X_test, y_test, k=11)
    
    # 7. 总结
    logger.info("="*60)
    logger.info("训练完成！")
    logger.info(f"索引名称: {INDEX_NAME}")
    logger.info(f"上传数据量: {len(X_train)} 条")
    logger.info(f"测试数据量: {len(X_test)} 条")
    logger.info(f"k=11 时的准确率: {accuracy:.4f}")
    logger.info("="*60)


if __name__ == "__main__":
    main()
